/*************************************************************************
 *                  HONGLING CAPITAL CONFIDENTIAL AND PROPRIETARY
 *
 *                COPYRIGHT (C) HONGLING CAPITAL CORPORATION 2012
 *    ALL RIGHTS RESERVED BY HONGLING CAPITAL CORPORATION. THIS PROGRAM
 * MUST BE USED  SOLELY FOR THE PURPOSE FOR WHICH IT WAS FURNISHED BY
 * HONGLING CAPITAL CORPORATION. NO PART OF THIS PROGRAM MAY BE REPRODUCED
 * OR DISCLOSED TO OTHERS,IN ANY FORM, WITHOUT THE PRIOR WRITTEN
 * PERMISSION OF HONGLING CAPITAL CORPORATION. USE OF COPYRIGHT NOTICE
 DOES NOT EVIDENCE PUBLICATION OF THE PROGRAM.
 *                  HONGLING CAPITAL CONFIDENTIAL AND PROPRIETARY
 *************************************************************************/

package com.springboot.common.aop;

import com.springboot.common.annotation.RateLimit;
import com.springboot.common.enums.LimitType;
import com.springboot.common.exception.RateLimitException;
import com.springboot.util.StringUtils;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Configuration;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.io.Serializable;
import java.lang.reflect.Method;
import java.util.Collections;
import java.util.List;

/**
 * 限流拦截器，用于限制方法的每秒请求次数，用来处理高并发问题。
 * <p>
 * <p>
 * 此处以IP为单位限流，可以考虑调整为全局。
 * </p>
 *
 * @author hongchao zhao at 2019-5-31 11:45
 */
@Aspect
@Configuration
public class LimitAspect {
    private static final Logger logger = LoggerFactory.getLogger(LimitAspect.class);

    @Autowired
    private RedisTemplate<String, Serializable> limitRedisTemplate;

    @Autowired
    private DefaultRedisScript<Number> redisluaScript;

    @Around("execution(* com.springboot.controller ..*(..) )")
    public Object interceptor(ProceedingJoinPoint joinPoint) throws Throwable {
        //获取拦截的方法名
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        //获取拦截的方法
        Method method = signature.getMethod();
        //获取方法名
        String functionName = method.getName();
        //获取拦截的方法类
        Class<?> targetClass = method.getDeclaringClass();
        //获取注解信息
        RateLimit rateLimit = method.getAnnotation(RateLimit.class);
        if (rateLimit != null) {
            //获取限流类型
            LimitType limitType = rateLimit.limitType();

            try {
                //生成唯一键值。
                String limitKey;
                switch (limitType) {
                    case IP:
                        limitKey = getIpAddress();
                        break;
                    case CUSTOMER:
                        limitKey = rateLimit.key();
                        break;
                    default:
                        limitKey = StringUtils.upperCase(method.getName());
                }

                //该键值决定了以什么规则来限制流速，经过测试限速不可取，也没必要，而且限速方法加进去后，每一次执行有误差
                String stringBuffer = targetClass.getName() + "-" + method.getName() + "-" + limitKey;
                List<String> keys = Collections.singletonList(stringBuffer);
                //执行Lua脚本，获取访问次数
                Number number = limitRedisTemplate.execute(redisluaScript, keys, rateLimit.count(), rateLimit.time());
                if (number != null && number.intValue() != 0 && number.intValue() <= rateLimit.count()) {
                    logger.info("限流时间段【{}s】内【{}】访问第：{} 次", rateLimit.time(), stringBuffer, number.toString());
                    return joinPoint.proceed();
                }
            } catch (Exception e) {
                throw new RateLimitException("当前请求【" + functionName + "】已经超出设置限流次数");
            }
        } else {
            //未添加当前限流注解的不作处理
            return joinPoint.proceed();
        }

        throw new RateLimitException("当前请求【" + functionName + "】已经超出设置限流次数");
    }

    /**
     * 获取当前请求IP。
     *
     * @return
     */
    public static String getIpAddress() {
        String ipAddress = null;
        try {
            HttpServletRequest request = ((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getRequest();
            ipAddress = request.getHeader("x-forwarded-for");
            if (ipAddress == null || ipAddress.length() == 0 || "unknown".equalsIgnoreCase(ipAddress)) {
                ipAddress = request.getHeader("Proxy-Client-IP");
            }
            if (ipAddress == null || ipAddress.length() == 0 || "unknown".equalsIgnoreCase(ipAddress)) {
                ipAddress = request.getHeader("WL-Proxy-Client-IP");
            }
            if (ipAddress == null || ipAddress.length() == 0 || "unknown".equalsIgnoreCase(ipAddress)) {
                ipAddress = request.getRemoteAddr();
            }
            // 对于通过多个代理的情况，第一个IP为客户端真实IP,多个IP按照','分割   "***.***.***.***".length()
            if (ipAddress != null && ipAddress.length() > 15) {
                // = 15
                if (ipAddress.indexOf(",") > 0) {
                    ipAddress = ipAddress.substring(0, ipAddress.indexOf(","));
                }
            }
        } catch (Exception e) {
            ipAddress = "";
        }
        return ipAddress;
    }
}
